import os
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from FSdata.FSdataset import FSdata, collate_fn
import torch
import torch.utils.data as torchdata
from models.resnet import resnet50_cat
from FSdata.FSaug import RandomHflip, Resize, Compose,Normalize
from utils.preprocessing import addFakeLabel
from utils.predicting import predict,gen_submission
import zipfile

class FSAugTest(object):
    def __init__(self):
        self.augment = Compose([
            Resize(size=(336,336),select=range(8)),
            Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def __call__(self, image, attr):
        return self.augment(image, attr)


os.environ["CUDA_VISIBLE_DEVICES"] = "2"

rawdata_root = '/media/gserver/data/FashionAI'

test_root = os.path.join(rawdata_root, 'rank')
test_pd = pd.read_csv(os.path.join(rawdata_root, 'rank/Tests/question.csv'),
                       header=None, names=['ImageName', 'AttrKey', 'AttrValues'])

test_pd = addFakeLabel(test_pd)


# test_root = os.path.join(rawdata_root, 'base')
# all_pd = pd.read_csv(os.path.join(rawdata_root, 'base/Annotations/label.csv'),
#                        header=None, names=['ImageName', 'AttrKey', 'AttrValues'])
#
# train_pd, test_pd = train_test_split(all_pd, test_size=0.1 ,random_state=37 ,
#                                     stratify=all_pd['AttrKey'])

print test_pd.shape

data_set = {}
data_set['test'] = FSdata(root_path=test_root,
                           anno_pd=test_pd,
                           transforms=FSAugTest(),
                          select=range(8)
                           )

data_loader = {}
data_loader['test'] = torchdata.DataLoader(data_set['test'], batch_size=16, num_workers=4,
                                          shuffle=False, pin_memory=True,collate_fn=collate_fn)


# model prepare
model_name = 'resnet50_cat_bs48_smooth'
resume = '/media/gserver/models/FashionAI/resnet50_cat_bs48_smooth/bestweights-[0.8732]-[0.9674].pth'
model = resnet50_cat(pretrained=True, num_classes=[8,5,5,5,10,6,6,9])

model = torch.nn.DataParallel(model)
print('resuming finetune from %s'%resume)
model.load_state_dict(torch.load(resume))
model = model.cuda()

# predict
if not os.path.exists('./online_pred/scores'):
    os.makedirs('./online_pred/scores')
if not os.path.exists('./online_pred/subs_csv'):
    os.makedirs('./online_pred/subs_csv')
if not os.path.exists('./online_pred/subs_zip'):
    os.makedirs('./online_pred/subs_zip')

# save csv and scores
pred_scores, _, _, _ = predict(model, data_set['test'], data_loader['test'])
test_pred = gen_submission(test_pd[['ImageName', 'AttrKey']], pred_scores)
test_pred[['ImageName', 'AttrKey', 'AttrValueProbs']].to_csv('online_pred/subs_csv/%s.csv'%model_name,header=None, index=False)
np.save('./online_pred/scores/%s.npy' % model_name,pred_scores)

# make zip file
z = zipfile.ZipFile('./online_pred/subs_zip/%s.zip'%model_name, 'w', zipfile.ZIP_DEFLATED)
z.write('./online_pred/subs_csv/%s.csv' % model_name, arcname='%s.csv' % model_name)
z.close()